#include "kernel_operator.h"

constexpr int32_t TOTAL_LENGTH = 2048;                                 // total length of data
constexpr int32_t USE_CORE_NUM = 8;                                   // num of core used
constexpr int32_t BLOCK_LENGTH = TOTAL_LENGTH / USE_CORE_NUM;         // length computed of each core
constexpr int32_t TILE_NUM = 8;                                       // split data into 8 tiles for each core
constexpr int32_t BUFFER_NUM = 2;                                     // tensor num for each queue
constexpr int32_t TILE_LENGTH = BLOCK_LENGTH / TILE_NUM / BUFFER_NUM; // separate to 2 parts, due to double buffer


using namespace AscendC;

template<typename T>
class KernelPowerInt {
public:
    __aicore__ inline KernelPowerInt() {}
    __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, GM_ADDR z, uint32_t totalLength, uint32_t tileNum)
    {
        this->blockLength = totalLength / GetBlockNum();
        this->tileNum = tileNum;
        this->tileLength = this->blockLength / tileNum / BUFFER_NUM;

        xGm.SetGlobalBuffer((__gm__ T *)x + this->blockLength * GetBlockIdx(), this->blockLength);
        yGm.SetGlobalBuffer((__gm__ int32_t *)y + this->blockLength * GetBlockIdx(), this->blockLength);
        zGm.SetGlobalBuffer((__gm__ T *)z + this->blockLength * GetBlockIdx(), this->blockLength);
        pipe.InitBuffer(inQueueX, BUFFER_NUM, this->tileLength * sizeof(T));
        pipe.InitBuffer(inQueueY, BUFFER_NUM, this->tileLength * sizeof(int32_t));
        pipe.InitBuffer(outQueueZ, BUFFER_NUM, this->tileLength * sizeof(T));

        
        pipe.InitBuffer(tmp1,  this->tileLength * sizeof(uint8_t));
        pipe.InitBuffer(tmp2,  this->tileLength * sizeof(T));
        pipe.InitBuffer(tmp3,  this->tileLength * sizeof(int32_t));
        pipe.InitBuffer(tmp4,  this->tileLength * sizeof(T));
        
    
        one_tmp= tmp2.Get<T>();
        Duplicate(one_tmp, (T)1, this->tileLength);

    }

    __aicore__ inline void Process()
    {
        int32_t loopCount = this->tileNum * BUFFER_NUM;
        for (int32_t i = 0; i < loopCount; i++) {
            CopyIn(i);
            Compute(i);
            CopyOut(i);
        }
    }

private:
    __aicore__ inline void CopyIn(int32_t progress)
    {
        LocalTensor<T> xLocal = inQueueX.AllocTensor<T>();
        LocalTensor<int32_t> yLocal = inQueueY.AllocTensor<int32_t>();
        DataCopy(xLocal, xGm[progress * this->tileLength], this->tileLength);
        DataCopy(yLocal, yGm[progress * this->tileLength], this->tileLength);
        inQueueX.EnQue(xLocal);
        inQueueY.EnQue(yLocal);
    }
    __aicore__ inline void Compute(int32_t progress)
    {
        LocalTensor<T> xLocal = inQueueX.DeQue<T>();
        LocalTensor<int32_t> yLocal = inQueueY.DeQue<int32_t>();
        LocalTensor<T> zLocal = outQueueZ.AllocTensor<T>();
        
        LocalTensor<uint8_t> bit= tmp1.Get<uint8_t>();
        LocalTensor<int32_t> tmpLocal= tmp3.Get<int32_t>();
        LocalTensor<T> repi= tmp4.Get<T>();

        Duplicate(zLocal, (T)1, this->tileLength);

        Div(repi,one_tmp, xLocal, this->tileLength);
       
        //CompareScalar(bit, yLocal, (int32_t)0, CMPMODE::LT, this->tileLength);
        //Select(xLocal, bit, repi, xLocal, SELMODE::VSEL_TENSOR_TENSOR_MODE, this->tileLength);
        //Abs(yLocal, yLocal, this->tileLength);
        for(int i=0;i<this->tileLength;i++){
            if(yLocal.GetValue(i)<0){
                yLocal.SetValue(i, -yLocal.GetValue(i));
                xLocal.SetValue(i, repi.GetValue(i));
            }
        }

        for(int i=0;i<31;i++){      //fast_power
            //And(tmpLocal,yLocal,and_tmp,this->tileLength);
            ShiftRight(tmpLocal, yLocal, (int32_t)1, this->tileLength);
            ShiftLeft(tmpLocal, tmpLocal, (int32_t)1, this->tileLength);
            Sub(tmpLocal, yLocal, tmpLocal, this->tileLength);
            //get the lowest bit in a slow way

            //Mul(repi, xLocal, zLocal, this->tileLength);
            //CompareScalar(bit, tmpLocal, (int32_t)0, CMPMODE::NE, this->tileLength);
            //Select(zLocal, bit, repi, zLocal, SELMODE::VSEL_TENSOR_TENSOR_MODE, this->tileLength);
            for(int i=0;i<this->tileLength;i++){
                if(tmpLocal.GetValue(i)!=0){
                    zLocal.SetValue(i, zLocal.GetValue(i)*xLocal.GetValue(i));
                }
            }
            ShiftRight(yLocal, yLocal, (int32_t)1, this->tileLength);
            Mul(xLocal, xLocal, xLocal, this->tileLength);
        }
      
        
        outQueueZ.EnQue<T>(zLocal);
        inQueueX.FreeTensor(xLocal);
        inQueueY.FreeTensor(yLocal);
    }
    __aicore__ inline void CopyOut(int32_t progress)
    {
        LocalTensor<T> zLocal = outQueueZ.DeQue<T>();
        DataCopy(zGm[progress * this->tileLength], zLocal, this->tileLength);
        outQueueZ.FreeTensor(zLocal);
    }

private:
    TPipe pipe;
    TQue<TPosition::VECIN, BUFFER_NUM> inQueueX, inQueueY;
    TQue<TPosition::VECOUT, BUFFER_NUM> outQueueZ;
    GlobalTensor<T> xGm;
    GlobalTensor<int32_t> yGm;
    GlobalTensor<T> zGm;

    LocalTensor<T> one_tmp;
    TBuf<QuePosition::VECCALC> tmp1,tmp2,tmp3,tmp4;

    uint32_t blockLength;
    uint32_t tileNum;
    uint32_t tileLength;
};

extern "C" __global__ __aicore__ void power_int_custom(GM_ADDR x, GM_ADDR y, GM_ADDR z)
{
    //GET_TILING_DATA(tiling_data, tiling);

    KernelPowerInt<float> op;
    op.Init(x, y, z, TOTAL_LENGTH,TILE_NUM);
    op.Process();
}



#ifndef ASCENDC_CPU_DEBUG
void power_int_custom_do(uint32_t blockDim, void *stream, float32_t *x, int32_t *y, float32_t *z)
{
    power_int_custom<<<blockDim, nullptr, stream>>>(x, y, z);
}
#endif
